/*
  mod.x - processes and modifies binary data
          compiling:  gcc mod_data.c -o mod.x -lm -lfftw3
*/
#include <tgmath.h>
#include <complex.h>
#include <stdio.h>
#include <fftw3.h>
#include <tgmath.h>
#include <string.h>
#include <stdlib.h>
#include <time.h>

static fftw_complex *rdata,*rdata2;
static unsigned int nx,ny,nz;
static double L,dx,dy,dz,ky,a;

void write_3d ( char* in ){
  FILE 			*fout;
  unsigned int 		i,j,k;
  double 		x[nx],y[ny],z[nz];
  
  fout = fopen(in,"w");
  
  fprintf(fout,"# File generated from binary input\n");   
  fprintf(fout,"# Full 3d domain\n"); 
  fprintf(fout,"# \tx\t\ty\t\tz\t\t|psi|\n\n");
  
  for (i = 0; i < nx; i++){
    for (j = 0; j < ny; j++){
      for (k = 0; k < nz; k++){
	x[i] = -0.5*L + (i+0.5)*dx;
	y[j] = -0.5*L + (j+0.5)*dy;
	z[k] = -0.5*L + (k+0.5)*dz;
	fprintf(fout,"%e\t%e\n",creal(rdata[k+nz*(j+ny*i)]), cimag(rdata[k+nz*(j+ny*i)]));
      }
      fprintf(fout,"\n");
    }
    fprintf(fout,"\n\n");
  }
  fclose(fout);
}

void write_3d_filtered ( char* in ){
  FILE 			*fout;
  unsigned int 		i,j,k;
  double 		x[nx],y[ny],z[nz];
  
  fout = fopen(in,"w");
  
  fprintf(fout,"# File generated from binary input\n");   
  fprintf(fout,"# Full 3d domain\n"); 
  fprintf(fout,"# \tx\t\ty\t\tz\t\t|psi|\n\n");
  
  for (i = 0; i < nx; i++){
    for (j = 0; j < ny; j++){
      for (k = 0; k < nz; k++){
	x[i] = -0.5*L + (i+0.5)*dx;
	y[j] = -0.5*L + (j+0.5)*dy;
	z[k] = -0.5*L + (k+0.5)*dz;
	fprintf(fout,"%e\t%e\t%e\t%e\n",x[i],y[j],z[k],cabs(rdata[k+nz*(j+ny*i)]));
      }
      fprintf(fout,"\n");
    }
    fprintf(fout,"\n\n");
  }
  fclose(fout);
}

void write_2d_slice(char* in, int num ){
  FILE 			*fout;
  unsigned int 		j,k;
  double 		y[ny],z[nz];
  
  fout = fopen(in,"w");
  
  fprintf(fout,"# File generated from binary input\n");   
  fprintf(fout,"# 2d slice yz-plane\n"); 
  fprintf(fout,"# \ty\t\tz\t\t|psi|\n\n");
  
    for (j = 0; j < ny; j++){
      for (k = 0; k < nz; k++){
	y[j] = -0.5*L + (j+0.5)*dy;
	z[k] = -0.5*L + (k+0.5)*dz;
	fprintf(fout,"%e\t%e\t%e\n",y[j],z[k],cabs(rdata[k+nz*(j+ny*num)]));
      }
      fprintf(fout,"\n");
    }
  fclose(fout);
}
void perturb_ky (char* in, double* kv, double a, int p_flag){
  //FILE 			*fout;
  unsigned int 		i,j,k;
  double 		x[nx],y[ny],z[nz], rr;
     
  for (i = 0; i < nx; i++){
    for (j = 0; j < ny; j++){
      for (k = 0; k < nz; k++){
	x[i] = -0.5*L + (i+0.5)*dx;
	y[j] = -0.5*L + (j+0.5)*dy;
	z[k] = -0.5*L + (k+0.5)*dz;
        rr = sqrt(x[i]*x[i]+y[j]*y[j]+z[k]*z[k]); 
	if (p_flag == 0) {
	  rdata[k+nz*(j+ny*i)] = rdata[k+nz*(j+ny*i)] + a*sin(kv[1]*x[i]+kv[2]*y[j]+kv[3]*z[k]);
	} else if (p_flag == 1) {
 	  rdata[k+nz*(j+ny*i)] = rdata[k+nz*(j+ny*i)]*exp(1I*(kv[1]*x[i]+kv[2]*y[j]+kv[3]*z[k]));
	} else if (p_flag == 2) {
	  rdata[k+nz*(j+ny*i)] = rdata[k+nz*(j+ny*i)]*exp( - pow(rr/(0.3*L),12) );
        }
      }
    }
  }

}

void move_sol ( double* r0) {
  unsigned int  ii, jj, kk;
  double	kkx, kky, kkz;
  fftw_plan	fplan, bplan;
  fftw_complex  *out, *out2;
  fftw_complex	DD;
  
  out = fftw_malloc(nx*ny*nz*sizeof(fftw_complex) );
  out2 = fftw_malloc(nx*ny*nz*sizeof(fftw_complex) );
  fplan = fftw_plan_dft_3d(nx, ny, nz, rdata, out, FFTW_FORWARD, FFTW_ESTIMATE );
  bplan = fftw_plan_dft_3d(nx, ny, nz, out, rdata, FFTW_BACKWARD, FFTW_ESTIMATE );
  fftw_execute(fplan);
  printf("in move_sol: %d\t%d\t%d\n", nx, ny, nz);
  printf("\t(r0x, r0y, r0z) = (%e,%e,%e)\n", r0[1], r0[2], r0[3]);

  for (ii = 0; ii < nx; ii++) for (jj = 0; jj < ny; jj++) for (kk = 0; kk < nz; kk++)
  {
    kkx = ii;
    kky = jj;
    kkz = kk;
    if (kkx >= nx/2) kkx = kkx - nx;
    if (kky >= ny/2) kky = kky - ny;
    if (kkz >= nz/2) kkz = kkz - nz;
    DD = cexp(2.I*M_PI*(1.*kkx*r0[1] + 1.*kky*r0[2] + 1.*kkz*r0[3]  )/L) /(nx*ny*nz);
    out2[kk+nz*(jj+ny*ii)] = out[kk+nz*(jj+ny*ii)]*DD;
    out[kk+nz*(jj+ny*ii)]  = out2[kk+nz*(jj+ny*ii)]; 
  }





  /*for (i = 0; i < nx; i++) {
   for (j = 0; j < ny; j++) {
    for (k = 0; k < nz; k++) {

    kx = 1.*i; 
    ky = 1.*j;
    kz = 1.*k;
    //if (ii > nx/2) kx = (kx - nx);
    //if (j > ny/2) ky = (ky - ny);
    //if (k > nz/2) kz = (kz - nz);
    mult = exp(1.I*(2.*M_PI)*(kx*r0[1] + ky*r0[2] + kz*r0[3])/L )/(nx*ny*nz);    
    out[k+nz*(j+ny*i)] = out[k+nz*(j+ny*i)]*mult;
    }
   }
  }*/
  fftw_execute(bplan);
}

void add_grids (){
  unsigned int 		i,j,k;
     
  for (i = 0; i < nx; i++){
    for (j = 0; j < ny; j++){
      for (k = 0; k < nz; k++){
	 rdata[k+nz*(j+ny*i)] = rdata[k+nz*(j+ny*i)] + rdata2[k+nz*(j+ny*i)];
      }
    }
  }

}



int main(int argc, char **argv ){
  
    FILE 		*fin, *fin2, *fh_mod;
    char		str[80];
    unsigned int	flag;
    double 		kv[4];
    double		r0[4];
    
    if ((argc != 8)&&(argc != 7)&&(argc != 11)) {
      printf("Usage:\n%s\tbin_file Flag [arg_1 arg_2 ...]\n", argv[0]);
      printf("Flag:\t0 converts binary output to text format:\n"); 
      printf("     \t  arg_1 - box length, arg_2, arg_3, arg_4 - Nx, Ny, Nz\n\n");	
      printf("     \t1 takes 2d yz-plane from binary output with z-index Nx/2 and converts to text format\n");
      printf("     \t  arg_1 - box length, arg_2, arg_3, arg_4 - Nx, Ny, Nz\n\n");	
      printf("     \t2 takes binary data and adds perturbation of the form:\n");
      printf("     \t\ta*sin((2*pi/L)*(kx*x+ky*y+kz*z)),\n     \t  where kx, ky, kz - arg_5, arg_6, arg_7 and a - arg_8\n");
      printf("     \t  and writes perturbed data to bin_file.mod file. arg_1, arg_2, arg_3, arg_4 are same as before.\n");
      printf("     \t  .mod file is in binary format for restarts.\n\n");	
      printf("     \t3 takes binary data and multiples by exp(i*(2*pi/L)*(kx*x+ky*y+kz*z))\n");
      printf("     \t  meaning of arguments same as for previous case, except arg_8, which is unused.\n");			
      printf("     \t  .mod file is in binary format for restarts.\n\n");
      printf("     \t4 takes two binary files and adds them together to .mod file\n");
      printf("     \t  arg_1 - box length, arg_2, arg_3, arg_4 - Nx, Ny, Nz, arg_5 - 2-nd binary file\n\n");
      printf("\t5 takes binary data and multiplies by a supergaussian with r_cutoff = 0.3L \n");	
      printf("     \t  arg_1 - box length, arg_2, arg_3, arg_4 - Nx, Ny, Nz\n\n");	
      printf("\t6 takes binary data and shifts it in direction r0, where r0x, r0y, r0z - arg_5, arg_6 and arg_7. arg_8 unused\n");
      printf("     \t  .mod file is in binary format for restarts.\n\n");
      printf("Example:\t mod.x binary_file 2 8 64 64 64 1 2 -1 0.1\n\n");		

      exit(1);
    }
    fin = fopen(argv[1],"rb");  
    if (fin == NULL) {
      printf("Can't open file\n");
      exit(1);
    }
   
    flag = atoi(argv[2]);  
    L    = atof(argv[3]);
    nx   = atoi(argv[4]); dx = L/nx;
    ny   = atoi(argv[5]); dy = L/ny;
    nz   = atoi(argv[6]); dz = L/nz;
    
    rdata  = malloc(nx*ny*nz*sizeof(fftw_complex));
	
    if (1 != fread(rdata,nx*ny*nz*sizeof(fftw_complex),1,fin)) {
	printf("Wrong size or corrupt data file\n");
	fclose(fin);
	exit(1);
    } else {
	printf("Successfully read data file\n");
    }    

    sprintf( str,"%s.dat",argv[1]);
    printf("Working with Flag: %d\n",flag);
    switch (flag) {
      case 0:

	   write_3d( str ); 
	   break;

      case 1:

	   write_2d_slice( str, nx/2 ); 
	   break;

      case 2: 

	   kv[1] = (2*M_PI/L)*atof(argv[7]) ;
	   kv[2] = (2*M_PI/L)*atof(argv[8]);
	   kv[3] = (2*M_PI/L)*atof(argv[9]);
	   perturb_ky( str, kv, atof(argv[10]), 0);    
	   sprintf( str,"%s.mod",argv[1]);
	   fh_mod = fopen(str,"wb");
	   fwrite(rdata, nx*ny*nz*sizeof(fftw_complex), 1, fh_mod);
	   fclose(fh_mod);
	   break;

      case 3:

	   kv[1] = (2*M_PI/L)*atof(argv[7]); 
	   kv[2] = (2*M_PI/L)*atof(argv[8]);
	   kv[3] = (2*M_PI/L)*atof(argv[9]);
	   perturb_ky( str, kv, 0, 1);    
	   sprintf( str,"%s.mod",argv[1]);
	   fh_mod = fopen(str,"wb");
	   fwrite(rdata, nx*ny*nz*sizeof(fftw_complex), 1, fh_mod);
	   fclose(fh_mod);
	   break;
     
      case 6:

	   r0[1] = atof(argv[7]);
	   r0[2] = atof(argv[8]);
	   r0[3] = atof(argv[9]);
  	   printf("Here %e\t%e\t%e\n", r0[1], r0[2], r0[3]);
           move_sol(r0);
	   sprintf( str,"%s.mod",argv[1]);
           fh_mod = fopen(str,"wb");
           fwrite(rdata, nx*ny*nz*sizeof(fftw_complex), 1, fh_mod);
           fclose(fh_mod);
           break;


      case 4:

	   rdata2 = malloc(nx*ny*nz*sizeof(fftw_complex));
	   fin2  = fopen(argv[7],"rb");  
	   if (1 == fread(rdata2,nx*ny*nz*sizeof(fftw_complex),1,fin2)) {	
            add_grids();
	     sprintf( str,"%s.mod",argv[1]);
	     fh_mod = fopen(str,"wb");
	     fwrite(rdata, nx*ny*nz*sizeof(fftw_complex), 1, fh_mod);
            fclose(fin2);
	     fclose(fh_mod);
	     break;
	   } else {
	     printf("Second file unreadable\n");
            fclose(fin2);
	     break;	
	   }

      case 5:

	   perturb_ky( str, 0, 0, 2);    
	   sprintf( str,"%s.mod",argv[1]);
	   fh_mod = fopen(str,"wb");
	   fwrite(rdata, nx*ny*nz*sizeof(fftw_complex), 1, fh_mod);
	   fclose(fh_mod);
	   break;



      default:

	   printf("Unknown flag.\n");
	   fclose(fin);
	   exit(1);
	   break;		

    }
   

    fclose(fin);
}


